"""腾讯混元大模型适配器实现。"""

from typing import AsyncGenerator, Dict, Any, List
import json
import hmac
import hashlib
import time
from datetime import datetime
import aiohttp

from .base import BaseModelAdapter, ModelConfig, ModelResponse


class HunyuanAdapter(BaseModelAdapter):
    """腾讯混元大模型适配器。"""
    
    # 定价（每千tokens，人民币）
    PRICING = {
        "hunyuan-standard": {"prompt": 0.005, "completion": 0.015},
        "hunyuan-lite": {"prompt": 0.002, "completion": 0.006},
        "hunyuan-pro": {"prompt": 0.03, "completion": 0.10}
    }
    
    def __init__(self, config: ModelConfig):
        """初始化混元适配器。
        
        需要在extra_params中提供:
        - secret_id: 腾讯云SecretId
        - secret_key: 腾讯云SecretKey
        - region: 地域，如ap-beijing
        """
        super().__init__(config)
        self.secret_id = config.extra_params.get("secret_id", "")
        self.secret_key = config.extra_params.get("secret_key", "")
        self.region = config.extra_params.get("region", "ap-beijing")
        
    def _initialize_client(self) -> None:
        """初始化混元客户端。"""
        # 腾讯云API地址
        self.base_url = self.config.api_base or f"https://hunyuan.tencentcloudapi.com"
    
    def _generate_auth_headers(self, payload: str) -> Dict[str, str]:
        """生成腾讯云API认证头。
        
        Args:
            payload: 请求体JSON字符串
            
        Returns:
            认证头字典
        """
        # 腾讯云API签名算法
        service = "hunyuan"
        host = "hunyuan.tencentcloudapi.com"
        algorithm = "TC3-HMAC-SHA256"
        timestamp = int(time.time())
        date = datetime.utcfromtimestamp(timestamp).strftime("%Y-%m-%d")
        
        # ************* 步骤 1：拼接规范请求串 *************
        http_request_method = "POST"
        canonical_uri = "/"
        canonical_querystring = ""
        canonical_headers = f"content-type:application/json\nhost:{host}\n"
        signed_headers = "content-type;host"
        hashed_request_payload = hashlib.sha256(payload.encode()).hexdigest()
        canonical_request = (
            f"{http_request_method}\n{canonical_uri}\n{canonical_querystring}\n"
            f"{canonical_headers}\n{signed_headers}\n{hashed_request_payload}"
        )
        
        # ************* 步骤 2：拼接待签名字符串 *************
        credential_scope = f"{date}/{service}/tc3_request"
        hashed_canonical_request = hashlib.sha256(canonical_request.encode()).hexdigest()
        string_to_sign = (
            f"{algorithm}\n{timestamp}\n{credential_scope}\n{hashed_canonical_request}"
        )
        
        # ************* 步骤 3：计算签名 *************
        def sign(key, msg):
            return hmac.new(key, msg.encode(), hashlib.sha256).digest()
        
        secret_date = sign(("TC3" + self.secret_key).encode(), date)
        secret_service = sign(secret_date, service)
        secret_signing = sign(secret_service, "tc3_request")
        signature = hmac.new(secret_signing, string_to_sign.encode(), hashlib.sha256).hexdigest()
        
        # ************* 步骤 4：拼接 Authorization *************
        authorization = (
            f"{algorithm} Credential={self.secret_id}/{credential_scope}, "
            f"SignedHeaders={signed_headers}, Signature={signature}"
        )
        
        return {
            "Authorization": authorization,
            "Content-Type": "application/json",
            "Host": host,
            "X-TC-Action": "ChatCompletions",
            "X-TC-Timestamp": str(timestamp),
            "X-TC-Version": "2023-09-01",
            "X-TC-Region": self.region
        }
    
    async def generate(
        self,
        messages: List[Dict[str, str]],
        **kwargs
    ) -> ModelResponse:
        """生成响应。
        
        Args:
            messages: 对话消息
            **kwargs: 额外参数
            
        Returns:
            模型响应
        """
        # 构建请求数据
        request_data = {
            "Model": self.config.model_name,
            "Messages": self._convert_messages(messages),
            "Temperature": self.config.temperature,
            "MaxTokens": self.config.max_tokens,
            "TopP": self.config.top_p,
            "Stream": False
        }
        
        # 添加额外参数
        if "enable_enhancement" in kwargs:
            request_data["EnableEnhancement"] = kwargs.pop("enable_enhancement")
        
        request_data.update(kwargs)
        
        # 生成请求体
        payload = json.dumps(request_data)
        
        # 生成认证头
        headers = self._generate_auth_headers(payload)
        
        async with aiohttp.ClientSession() as session:
            async with session.post(
                self.base_url,
                headers=headers,
                data=payload,
                timeout=aiohttp.ClientTimeout(total=self.config.timeout)
            ) as response:
                result = await response.json()
                
                if "Error" in result.get("Response", {}):
                    error = result["Response"]["Error"]
                    raise Exception(f"混元API错误: {error.get('Message', 'Unknown error')}")
                
                response_data = result.get("Response", {})
                choices = response_data.get("Choices", [])
                
                if not choices:
                    raise Exception("混元API返回空响应")
                
                choice = choices[0]
                usage = response_data.get("Usage", {})
                
                return ModelResponse(
                    content=choice.get("Message", {}).get("Content", ""),
                    model=self.config.model_name,
                    usage={
                        "prompt_tokens": usage.get("PromptTokens", 0),
                        "completion_tokens": usage.get("CompletionTokens", 0),
                        "total_tokens": usage.get("TotalTokens", 0)
                    },
                    finish_reason=choice.get("FinishReason", "stop"),
                    metadata={
                        "request_id": response_data.get("RequestId", ""),
                        "created": response_data.get("Created", 0)
                    }
                )
    
    async def stream(
        self,
        messages: List[Dict[str, str]],
        **kwargs
    ) -> AsyncGenerator[str, None]:
        """流式生成响应。
        
        注意：腾讯云SDK的流式实现较复杂，这里提供简化版本。
        
        Args:
            messages: 对话消息
            **kwargs: 额外参数
            
        Yields:
            内容块
        """
        # 混元流式API实现
        request_data = {
            "Model": self.config.model_name,
            "Messages": self._convert_messages(messages),
            "Temperature": self.config.temperature,
            "MaxTokens": self.config.max_tokens,
            "TopP": self.config.top_p,
            "Stream": True
        }
        
        request_data.update(kwargs)
        payload = json.dumps(request_data)
        headers = self._generate_auth_headers(payload)
        
        # 注意：实际流式实现需要处理SSE格式
        async with aiohttp.ClientSession() as session:
            async with session.post(
                self.base_url,
                headers=headers,
                data=payload,
                timeout=aiohttp.ClientTimeout(total=self.config.timeout)
            ) as response:
                async for line in response.content:
                    if line:
                        # 简化处理，实际需要解析SSE格式
                        yield "混元流式响应处理中..."
                        break
    
    def _convert_messages(self, messages: List[Dict[str, str]]) -> List[Dict[str, str]]:
        """转换消息格式以适配混元API。
        
        Args:
            messages: 原始消息列表
            
        Returns:
            转换后的消息列表
        """
        converted = []
        for msg in messages:
            role = msg["role"]
            # 混元使用的角色名称
            if role == "system":
                role = "system"
            elif role == "assistant":
                role = "assistant"
            else:
                role = "user"
            
            converted.append({
                "Role": role,
                "Content": msg["content"]
            })
        
        return converted
    
    def count_tokens(self, text: str) -> int:
        """计算文本的token数量。
        
        混元的token计算方式。
        
        Args:
            text: 输入文本
            
        Returns:
            Token数量
        """
        # 中文字符约1个token
        chinese_chars = sum(1 for char in text if '\u4e00' <= char <= '\u9fff')
        # 英文约4个字符1个token
        english_chars = sum(1 for char in text if char.isascii())
        
        return chinese_chars + (english_chars // 4)
    
    def estimate_cost(self, prompt_tokens: int, completion_tokens: int) -> float:
        """估算使用成本。
        
        Args:
            prompt_tokens: 提示词token数
            completion_tokens: 生成token数
            
        Returns:
            预估成本（美元）
        """
        # 根据模型名称判断定价层级
        if "pro" in self.config.model_name.lower():
            model_key = "hunyuan-pro"
        elif "lite" in self.config.model_name.lower():
            model_key = "hunyuan-lite"
        else:
            model_key = "hunyuan-standard"
        
        pricing = self.PRICING[model_key]
        prompt_cost = (prompt_tokens / 1000) * pricing["prompt"]
        completion_cost = (completion_tokens / 1000) * pricing["completion"]
        
        # 转换为美元
        cny_to_usd = 0.14
        return (prompt_cost + completion_cost) * cny_to_usd